# 作者 :南雨
# 时间 : 2022/6/28 9:12
import nltk
from nltk.corpus import stopwords
from dzj.trec_qa.trec_data_processing import get_all
from nltk.corpus import wordnet
from nltk.stem import WordNetLemmatizer
from dzj.trec_qa.trec_constant import Trec_constant
import re

baidu_stopwords = [line.strip() for line in open(Trec_constant.baidu_stop_words, encoding="utf-8").readlines()]  # 加载停用词


def get_wordnet_pos(tag):
    if tag.startswith('J'):
        return wordnet.ADJ
    elif tag.startswith('V'):
        return wordnet.VERB
    elif tag.startswith('N'):
        return wordnet.NOUN
    elif tag.startswith('R'):
        return wordnet.ADV
    else:
        return None


def process(text):
    """
    分词、去停用词、小写转换
    :param text:一行文本
    :return:["word1","word2"]
    """
    text01 = text.lower()
    result_list = []
    wnl = WordNetLemmatizer()
    # stop_words = stopwords.words('english')
    for filters in ['\'', '!', ',', '.', '?', '-s', '-ly', '</s>', 's', '\'s', '\'\'']:
        baidu_stopwords.append(filters)
    filters = ['!', '"', '#', '$', '%', '&', '\(', '\)', '\*', '\+', ',', '-', '\.', '/', ':', ';', '<', '=', '>',
               '\?', '@', '\[', '\\', '\]', '^', '_', '`', '\{', '\|', '\}', '~', '\t', '\n', '\x97', '\x96', '”',
               '“', ]
    text02 = re.sub("|".join(filters), "", text01).strip()
    text03 = re.sub("\s+", "*", text02).strip()
    text_list = text03.split("*")
    # 词干还原
    tag_text = nltk.pos_tag(text_list)
    for tag_word in tag_text:
        word_ori = tag_word[0]
        word_tag = tag_word[1]
        word_tag = get_wordnet_pos(word_tag) or wordnet.NOUN
        word = wnl.lemmatize(word_ori, word_tag)
        # if word not in stop_words:
        result_list.append(word)
    return result_list


def get_result1(text):
    result_list = process(text)
    return ' '.join(result_list)


def get_process_df(df):
    df['text'] = df['text'].apply(get_result1)
    return df


def get_result_list(text_list):
    word_list = []
    for i in text_list:
        res = process(i[0])
        word_list.append(res)
    return word_list


def result():
    data_df, text_list = get_all()
    return get_result_list(text_list), get_process_df(data_df)


if __name__ == '__main__':
    print(result()[0])
